'''
将文件夹下的dcm文件保存为nii.gz文件
'''
import os
import time
import SimpleITK

import shutil

def dicom_to_nifti(series_dir, give_name):

    # https://simpleitk.readthedocs.io/en/master/link_DicomConvert_docs.html
    if os.path.isdir(series_dir) and len(os.listdir(series_dir)) > 1:
        reader = SimpleITK.ImageSeriesReader()
        dicom_names = reader.GetGDCMSeriesFileNames(series_dir)
        reader.SetFileNames(dicom_names)
        image = reader.Execute()
    else:
        filename = (
            series_dir if not os.path.isdir(series_dir) else os.path.join(series_dir, os.listdir(series_dir)[0])
        )

        file_reader = SimpleITK.ImageFileReader()
        file_reader.SetImageIO("GDCMImageIO")
        file_reader.SetFileName(filename)
        image = file_reader.Execute()
    #output_file = tempfile.NamedTemporaryFile(suffix=".nii.gz").name
    output_file = give_name
    SimpleITK.WriteImage(image, give_name)

    return output_file

def copySegmentation(source_folder_path,segmentation_save_path):
    if not os.path.exists(segmentation_save_path):
        print("create directory {}".format(segmentation_save_path))
        os.makedirs(segmentation_save_path)
    for file_name in os.listdir(source_folder_path):
        file_directory = os.path.join(source_folder_path, file_name)
        seg_name=file_name
        for file in os.listdir(file_directory):
            if '.nii' in file:
                nii_path=os.path.join(file_directory,file)
                copy_path=os.path.join(segmentation_save_path,seg_name+'.nii.gz')
                print("copy file {} to {}".format(nii_path, copy_path))
                shutil.copyfile(nii_path,copy_path)

if __name__ == '__main__':
    volume_path = "E:/AllData/zhongShanBreast/nii/volume"  # DCM文件转化为nii格式后的路径
    sus_ROI_path = "E:/AllData/zhongShanBreast/sus_ROI"  # DCM文件路径
    normalExamples_path = "E:/AllData/zhongShanBreast/normalExamples"  # normalExamples的DCM文件路径
    seg_path = "E:/AllData/zhongShanBreast/nii/segmentation"
    if not os.path.exists(volume_path):
        print("create directory {}".format(volume_path))
        os.makedirs(volume_path)

    for file_name in os.listdir(sus_ROI_path):
        file_path = os.path.join(sus_ROI_path,file_name)
        for i in os.listdir(file_path):
            if 'c' in i:
                end_folder_path = os.path.join(file_path,i)
                print(end_folder_path)
                give_name = os.path.join(volume_path , file_name+".nii.gz")
                print(give_name)
                print("******",end_folder_path,give_name)
                dicom_to_nifti(end_folder_path, give_name)
    # #下面处理正常病例的
    for file_name in os.listdir(normalExamples_path):
        if 'c' in file_name:
            end_folder_path = os.path.join(normalExamples_path,file_name)
            print(end_folder_path)
            give_name = os.path.join(volume_path , file_name+".nii.gz")
        print(give_name)
        dicom_to_nifti(end_folder_path, give_name)
    copySegmentation('E:\AllData\zhongShanBreast\sus_ROI',"E:/AllData/zhongShanBreast/nii/segmentation")